package astudy.bytecode.bcel;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;

import org.apache.bcel.Constants;
import org.apache.bcel.classfile.ClassParser;
import org.apache.bcel.classfile.JavaClass;
import org.apache.bcel.classfile.Method;
import org.apache.bcel.generic.ClassGen;
import org.apache.bcel.generic.ConstantPoolGen;
import org.apache.bcel.generic.InstructionConstants;
import org.apache.bcel.generic.InstructionFactory;
import org.apache.bcel.generic.InstructionList;
import org.apache.bcel.generic.MethodGen;
import org.apache.bcel.generic.ObjectType;
import org.apache.bcel.generic.PUSH;
import org.apache.bcel.generic.Type;
import org.apache.commons.io.IOUtils;

import astudy.bytecode.ClassFileWrapper;
import astudy.bytecode.NamePattern;

public class ProfilerWrapper implements ClassFileWrapper {
	
    public byte[] wrap(byte[] bytes, NamePattern methodNamePattern) 
    throws IOException {
       
    	//create temporary file
    	File classFile = File.createTempFile("jtmp_", ".class");
    	OutputStream fos = new FileOutputStream(classFile); 
    	IOUtils.write(bytes, fos);
    	IOUtils.closeQuietly(fos);
    	//modify the class file
    	String classFilepath = classFile.getAbsolutePath();
        JavaClass jclas = new ClassParser(classFilepath).parse();
        ClassGen cgen = new ClassGen(jclas);
        
        for (Method method : jclas.getMethods()) {
            if (methodNamePattern.match(method.getName())) {
                addWrapper(cgen, method);
            }
        }
        //dump
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        cgen.getJavaClass().dump(bos);
        bos.close();
        return bos.toByteArray();
    }

    private void addWrapper(ClassGen cgen, Method method) {
        
        // set up the construction tools
        InstructionFactory ifact = new InstructionFactory(cgen);
        ConstantPoolGen pgen = cgen.getConstantPool();
        String cname = cgen.getClassName();
        MethodGen wrapgen = new MethodGen(method, cname, pgen);
        InstructionList ilist = new InstructionList();
        wrapgen.setInstructionList(ilist);
        
        // rename a copy of the original method
        MethodGen methgen = new MethodGen(method, cname, pgen);
        cgen.removeMethod(method);
        String iname = methgen.getName() + "$impl";
        methgen.setName(iname);
        cgen.addMethod(methgen.getMethod());
        Type result = methgen.getReturnType();
        
        // save time prior to invocation
        //
        // invokestatic java.lang.System.currentTimeMillis
        //
        ilist.append(ifact.createInvoke("java.lang.System",
            "currentTimeMillis", Type.LONG, Type.NO_ARGS, 
            Constants.INVOKESTATIC));
        
        // compute the size of the calling parameters
        //
        // lstore_2
        //
        Type[] types = methgen.getArgumentTypes();
        int slot = methgen.isStatic() ? 0 : 1;
        for (int i = 0; i < types.length; i++) {
            slot += types[i].getSize();
        }
        
        ilist.append(InstructionFactory.createStore(Type.LONG, slot));
        
        // call the wrapped method
        //
        // aload_0
        // iload_1
        //
        int offset = 0;
        short invoke = Constants.INVOKESTATIC;
        if (!methgen.isStatic()) {
            ilist.append(InstructionFactory.createLoad(Type.OBJECT, 0));
            offset = 1;
            invoke = Constants.INVOKEVIRTUAL;
        }
        for (int i = 0; i < types.length; i++) {
            Type type = types[i];
            ilist.append(InstructionFactory.createLoad(type, offset));
            offset += type.getSize();
        }
        //
        // invokevirtual xxx.buildString$impl 
        //
        ilist.append(ifact.createInvoke(cname, 
            iname, result, types, invoke));
        
        // store result for return later
        //
        // astore 4
        //
        if (result != Type.VOID) {
            ilist.append(InstructionFactory.
                createStore(result, slot+2));
        }
        
        // print time required for method call
        //
        // getstatic java.lang.System.out
        // dup
        // dup
        //
        ilist.append(ifact.createFieldAccess("java.lang.System",
            "out",  new ObjectType("java.io.PrintStream"),
            Constants.GETSTATIC));
        ilist.append(InstructionConstants.DUP);
        ilist.append(InstructionConstants.DUP);
        //
        // ldc "Call to method ... took"
        // invokevirtual java.io.PrintStream.print
        //
        String text = "Call method : " + method.getName() + "() - ";
        ilist.append(new PUSH(pgen, text));
        ilist.append(ifact.createInvoke("java.io.PrintStream",
            "print", Type.VOID, new Type[] { Type.STRING },
            Constants.INVOKEVIRTUAL));
        //
        // invokestatic java.lang.System.currentTimeMillis
        // lload_2
        // lsub
        // invokevirtual java.io.PrintStream.print
        // ldc "ms."
        // invokevirtual java.io.PrintStream.print
        //
        ilist.append(ifact.createInvoke("java.lang.System", 
            "currentTimeMillis", Type.LONG, Type.NO_ARGS, 
            Constants.INVOKESTATIC));
        ilist.append(InstructionFactory.createLoad(Type.LONG, slot));
        ilist.append(InstructionConstants.LSUB);
        ilist.append(ifact.createInvoke("java.io.PrintStream",
            "print", Type.VOID, new Type[] { Type.LONG },
            Constants.INVOKEVIRTUAL));
        ilist.append(new PUSH(pgen, " ms."));
        ilist.append(ifact.createInvoke("java.io.PrintStream",
            "println", Type.VOID, new Type[] { Type.STRING },
            Constants.INVOKEVIRTUAL));
            
        // return result from wrapped method call
        //
        // aload 4
        //
        if (result != Type.VOID) {
            ilist.append(InstructionFactory.
                createLoad(result, slot+2));
        }
        ilist.append(InstructionFactory.createReturn(result));
        
        // finalize the constructed method
        wrapgen.stripAttributes(true);
        wrapgen.setMaxStack();
        wrapgen.setMaxLocals();
        cgen.addMethod(wrapgen.getMethod());
        ilist.dispose();
    }
} 
